package spark.sql

import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types._

/**
 * Created by ibf on 2018/2/10.
 */
object Avg_UDAF extends UserDefinedAggregateFunction{

  /**
   * 指定自定义的UDAF函数的输入的数据类型
   * @return  返回值是一个StructType，这个值可以封装多个字段
   */
  override def inputSchema: StructType = {
    StructType(Array(
      StructField("iv",DoubleType)
    ))
  }

  /**
   * 如何去更新buffer缓存区中的值
   * @param buffer
   * @param input
   */
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    //对于每一条输入的数据（相同分组），更新buffer中的值
    //1、获取输入的值
    val iv = input.getDouble(0)
    //2、获取缓存区中的值
    //tv是上一次聚合的薪水，tc是上一次聚合的薪水对应的人数
    val tv = buffer.getDouble(0)
    val tc = buffer.getInt(1)

    //3、将iv的值，累加到tv上，更新buffer的数据
    buffer.update(0,tv + iv)
    buffer.update(1,tc + 1)
  }

  /**
   * 聚合函数
   * 多条数据输入，返回一条数据
   * 多条数据进行计算，有一个buffer来存储上一次计算的值
   * 所以要指定缓存区的schema信息
   * @return
   *
   */
  override def bufferSchema: StructType = {
    //两个值，第一个是上一次的薪水的累加值，第二个是上一次的人数累加值
    StructType(Array(
      StructField("tv",DoubleType),
      StructField("tc",IntegerType)
    ))
  }

  /**
   * 指定相同分区的多个buffer值，如何进行累加
   * @param buffer1 第一个buffer1，是MutableAggregationBuffer类型
   * @param buffer2 第二个buffer2，是Row类型
   */
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    //1、获取buffer1中的数据
    val tv1 = buffer1.getDouble(0)
    val tc1 = buffer1.getInt(1)
    //2、获取buffer2中的数据
    val tv2 = buffer2.getDouble(0)
    val tc2 = buffer2.getInt(1)

    /**
     * 3、把缓存区中的数据进行更新
     * 问题：把buffer1中的数据更新到buffer2中？
     *      把buffer2中的数据更新到buffer1中？
     * 答案：因为MutableAggregationBuffer类型提供了update的接口
     *      而普通的Row是没有更新方法的，所以buffer2更新到buffer1上
     */
    buffer1.update(0,tv1+tv2)
    buffer1.update(1,tc1+tc2)
  }

  /**
   * 指定构造出来的缓存区的初始值
   * @param buffer
   */
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    //薪水初始值为0.0，因为薪水是double类型的
    buffer.update(0,0.0)
    //统计的人数，给定0，因为人数的int类型
    buffer.update(1,0)
  }

  /**
   * 给定多次运行该方法，给定相同的输入值，是否返回相同的结果
   * 默认是true表示允许
   * @return
   */
  override def deterministic: Boolean = {
    true
  }

  /**
   * evaluate
   * hive 自定义函数，evaluate来写运算逻辑的方法名称
   * 所以spark中也是一样，在这里实现运算逻辑
   * @param buffer  全部数据都累加完之后的缓存
   * @return  想要返回的数据，在当前案例中，是返回平均薪水
   */
  override def evaluate(buffer: Row): Any = {
    //1、获取所有数据累加完之后的薪水值（当前分区）
    val tv = buffer.getDouble(0)
    //2、获取当前分区（部门）的总人数
    val tc = buffer.getInt(1)
    //3、得到平均值，并且返回平均值
      tv / tc
  }

  /**
   * 指定返回的数据类型
   * @return DataType类型
   */
  override def dataType: DataType = {
    DoubleType
  }
}
